import traceback
from django.conf import settings
from channels.generic.websocket import AsyncJsonWebsocketConsumer
from django.contrib.auth.models import AnonymousUser
from channels.db import database_sync_to_async
from ws.groups import DATA_QUERY_GROUPS
from channels.generic.websocket import SyncConsumer
from asgiref.sync import async_to_sync


class AsyncWSConsumer(AsyncJsonWebsocketConsumer):
	def __init__(self, *args, **kwargs):
		super().__init__(*args, **kwargs)
		self.group_kwargs = {}

	async def connect(self):
		"""ws token 认证连接"""
		await super().connect()
		user = self.scope.get('user')
		if user is None or isinstance(user, AnonymousUser):
			await self.send_json({'action': 'fail', 'message': 'Signature verification failed'})
			await self.close()

	async def receive_json(self, content, **kwargs):
		"""
		1.处理前端心跳
		2.option 功能分类
		3.必传字段报警处理
		"""
		# 前端心跳响应
		if "heartbeat" in content:
			await self.send_json({"heartbeat": "heartbeat"})
		option = content.get('option')		# 选项、类型 用于区分使用自定义方法
		group = content.get('group')		# 分组
		if not all([option, group]):
			await self.send_json(
				{
					"action": "error",
					"message": "No 'option' or 'group' input",
					"data": {}
				}
			)

		self.group_name = group
		self.group_kwargs[group] = content		# PeriodicConsumer 中使用
		option_list = ['subscribe', 'unsubscribe', 'request']
		if option in option_list:
			func = getattr(self, option)
			await func(group, content)
		else:
			await self.send_json(
				{
					"action": "error",
					"message": "Not valid option",
					"data": {}
				}
			)

	async def subscribe(self, group, content):
		"""实现订阅"""
		res = await self.get_data(group, content)
		# 实现广播、组织连接：使用通道层将连接加入到组中，以便可以向组中的所有连接发送消息
		await self.channel_layer.group_add(group, self.channel_name)
		await self.send_json(res)

	async def unsubscribe(self, group, content):
		"""取消订阅"""
		await self.channel_layer.group_discard(group, self.channel_name)
		await self.send_json(
			{
				"action": "success",
				"message": f"取消订阅：{group}",
				"data": {}
			}
		)

	async def request(self, group, content):
		"""处理请求,非订阅"""
		res = await self.get_data(group, content)
		await self.send_json(res)

	@database_sync_to_async
	def get_data(self, group, content):
		"""映射分发"""
		func = DATA_QUERY_GROUPS.get(group)
		if func:
			try:
				data = func(content)
				error_msg = ''
			except:
				data = None
				error_msg = traceback.format_exc()
				if settings.DEBUG is True:
					traceback.print_exc()
			res = {"action": "success", "message": error_msg, "data": data}
		else:
			res = {"action": "error", "message": f"分组：{group} 不存在", "data": None}
		return res

	async def periodic_sending(self, content):
		message = content.get("message")
		group_name = message.get("group")

		kwargs = self.group_kwargs.get(group_name)
		res = await self.get_data(group_name, **kwargs)
		await self.send_json(res)


class PeriodicConsumer(SyncConsumer):
	def broadcast(self, content):
		message = content.get("message")
		group = message.get("group")
		async_to_sync(self.channel_layer.group_send)(
			group,
			{
				# type 为对应 consumer 类中的消费函数
				"type": "periodic_sending",
				"message": message,
			}
		)
